#lang racket

(require "util/utils.scm")
(provide pass-tag)

;; inputs:
;;   source is a list of top-level raw s-expressions
;;
;; outputs:
;;   the output will be a list of tagged forms that have been syntax checked
;;   variables like x are now (var x), applications (f x) are now (app f x)
;;
;;   list of the tags appearing in output (search ` to make this list):
;;
;;   toplevel:
;;   * define
;;   * dont-define
;;
;;   expressions:
;;   * atomic
;;   * var
;;   * app
;;   * begin
;;   * lambda
;;   * quote
;;   * if
;;   * let*

(define special-forms '(lambda begin quote if let* letrec builtin))

(define (pass-tag source) (map* pass-tag-toplevel source))

;; This handles the top level definitions

(define (pass-tag-toplevel def)
  (if (pair? def)
      (match def
        (`(define ,head . ,body)
         (cond ((symbol? head)
                (unless (null? (cdr body))
                  (error 'pass-tag-toplevel "multiple expressions in the definition of a variable" def))
                `(define ,head ,(pass-tag-expr (car body))))
               ((pair? head)
                (pass-tag-toplevel `(define ,(car head) (lambda ,(cdr head) . ,body))))
               (else
                (error 'pass-tag-toplevel "bad head for define" def))))
        (else `(dont-define ,(pass-tag-expr def))))
      `(dont-define ,(pass-tag-expr def))))


;; This handles scheme expressions

(define (pass-tag-expr exp)
  (pass-tag-expr* exp special-forms))

(define (pass-tag-expr* exp forms)
  (cond ((atomic? exp) `(atomic ,exp))
        ((symbol? exp) `(var ,exp))
        ((pair? exp)
         (let ((h (car exp))
               (t (cdr exp)))
           (if (member h forms)
               (case h
                 ((lambda) (pass-tag-lambda forms t))
                 ((begin) (pass-tag-begin forms t))
                 ((quote) (pass-tag-quote forms t))
                 ((if) (pass-tag-if forms t))
                 ((let*) (pass-tag-let* forms t))
                 ((letrec) (pass-tag-letrec forms t))
                 ((builtin) (pass-tag-builtin forms t))
                 (else (error 'pass-tag-expr* "impossible situation" (list h forms))))
               `(app . ,(mapply pass-tag-expr* exp forms)))))
        (else (error 'pass-tag-expr* "unknown object" exp))))


;; These are validation and tagging methods for specific special forms

(define (pass-tag-lambda forms t)
  (unless (and (pair? t) (pair? (cdr t)))
    (error 'pass-tag-lambda "malformed lambda" `(lambda . ,t)))
  (let ((bindings (car t)))
    (unless (list? bindings)
      (error 'pass-tag-lambda "lambda bindings do not form a list" bindings))
    (let* ((forms (removes bindings forms))
           (body (pass-tag-implicit-begin forms (cdr t))))
      `(lambda ,bindings ,body))))

(define (pass-tag-begin forms t)
 `(begin . ,(mapply pass-tag-expr* t forms)))

(define (pass-tag-quote forms t)
  (unless (null? (cdr t))
    (error 'pass-tag-quote "quote only accepts a single expression" `(quote . ,t)))
 `(quote ,(car t)))

(define (pass-tag-if forms t)
  (unless (length=? t 3)
    (error 'pass-tag-if "if needs 3 parameters" `(if . ,t)))
  `(if ,(pass-tag-expr* (car t) forms)
       ,(pass-tag-expr* (cadr t) forms)
       ,(pass-tag-expr* (caddr t) forms)))

(define (pass-tag-let* forms t)
  (unless (length>=? t 2)
    (error 'pass-tag-let* "let* needs a bindings list then body forms" `(let* . ,t)))
  (for-each (lambda (bind)
              (unless (and (length=? bind 2) (symbol? (car bind)))
                (error 'pass-tag-let* "let* contained a binding not of the form (<var> <exp>)" bind)))
            (car t))
  (let* ((bindings (let loop ((binds (car t)))
                     (if (null? binds)
                         '()
                         (let ((tmp (caar binds))
                               (exp (pass-tag-expr* (cadar binds) forms)))
                           (set! forms (remove tmp forms))
                           (cons (list tmp exp) (loop (cdr binds)))))))
         (body (pass-tag-implicit-begin forms (cdr t))))
    `(let* ,bindings ,body)))

(define (pass-tag-letrec forms t)
  ;; can be merged with let*
  (unless (length>=? t 2)
    (error 'pass-tag-let* "letrec needs a bindings list then body forms" `(let* . ,t)))
  (for-each (lambda (bind)
              (unless (and (length=? bind 2) (symbol? (car bind)))
                (error 'pass-tag-let* "letrec contained a binding not of the form (<var> <exp>)" bind)))
            (car t))
  (let* ((bindings (let loop ((binds (car t)))
                     (if (null? binds)
                         '()
                         (let ((tmp (caar binds))
                               (exp (pass-tag-expr* (cadar binds) forms)))
                           (set! forms (remove tmp forms))
                           (cons (list tmp exp) (loop (cdr binds)))))))
         (body (pass-tag-implicit-begin forms (cdr t))))
    `(letrec ,bindings ,body)))

(define (pass-tag-builtin forms t)
  (unless (and (length>=? t 1) (symbol? (car t)))
    (error 'pass-tag-builtin "malformed builtin call" `(builtin . ,t)))
  `(builtin ,(car t) . ,(mapply pass-tag-expr* (cdr t) forms)))


;; utility function for adding an implement begin when needed

(define (pass-tag-implicit-begin forms body)
  (let ((single (null? (cdr body))))
    (if single
        (pass-tag-expr* (car body) forms)
        `(begin . ,(mapply pass-tag-expr* body forms)))))
